package com.eobard.obfuscate.core;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import com.eobard.adapter.BaseVisitorAdapter;
import com.eobard.config.obfuscate.ObfuscatorConfig;
import com.eobard.enums.RunningEnum;
import com.eobard.obfuscate.IObfuscator;
import com.eobard.params.IMethod;
import com.eobard.params.MethodArgumentParams;
import com.eobard.utils.CollectionUtils;
import com.eobard.utils.ExceptionUtils;
import com.eobard.utils.ObfuscatorUtils;
import com.eobard.utils.RandomExpression;
import com.github.javaparser.StaticJavaParser;
import com.github.javaparser.ast.CompilationUnit;
import com.github.javaparser.ast.Modifier;
import com.github.javaparser.ast.NodeList;
import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
import com.github.javaparser.ast.body.FieldDeclaration;
import com.github.javaparser.ast.body.MethodDeclaration;
import com.github.javaparser.ast.body.Parameter;
import com.github.javaparser.ast.expr.Expression;
import com.github.javaparser.ast.expr.MethodCallExpr;
import com.github.javaparser.ast.stmt.Statement;
import com.github.javaparser.ast.type.Type;
import com.github.javaparser.resolution.declarations.ResolvedMethodDeclaration;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * @Author: divenswu(核心代码)、Eobard Gu(重构代码)
 * @Date: 2025-05-22
 * @Description: 修改普通方法参数核心混淆类
 */
@Component
public class MethodArgumentObfuscator extends BaseVisitorAdapter<MethodArgumentParams> implements IObfuscator {

    @Value("${obfuscate.enable.typeSolver.logging:false}")
    private boolean print;
    private final Random random = new Random();

    private String sourceRoot;
    private final Map<String, CompilationUnit> fileMap = new HashMap<>();
    private static List<MethodDeclaration> tempStaticMethods = new ArrayList<>();
    private final Map<String, List<MethodCallExpr>> methodCallExprMap = new ConcurrentHashMap<>();
    private final Map<String, MethodDeclaration> methodDeclarationMap = new ConcurrentHashMap<>();


    //数据初始化代码
    @Override
    public void initialize(ObfuscatorConfig config) {
        if (StrUtil.isBlank(config.getSourceRoot())) throw new RuntimeException("请先设置项目路径!");
        //清空上一次记录
        CollectionUtils.clear(fileMap, methodCallExprMap, methodDeclarationMap);
        CollectionUtils.clear(tempStaticMethods);

        ExceptionUtils.checkNonNull(config.getMethodArgumentParams().getBasePackage(), "请先配置【普通方法参数混淆器】的basePackage!");
        ExceptionUtils.checkNonNull(config.getMethodArgumentParams().getPackageList(), "请先配置【普通方法参数混淆器】的包列表!");

        this.sourceRoot = config.getSourceRoot();
        initSolver(sourceRoot, config.getSourceExt());

        //获取源码文件
        ObfuscatorUtils.loopFiles(fileMap, sourceRoot);
        //获取模版里的静态函数
        tempStaticMethods = ObfuscatorUtils.getStaticTemplateMethods(staticMethodsFilePath);
    }

    //核心混淆代码
    @Override
    public void obfuscate(ObfuscatorConfig config) {
        MethodArgumentParams argumentParams = config.getMethodArgumentParams();
        //计算当前需要修改的函数方法
        this.fileMap.forEach((path, cu) -> cu.accept(this, argumentParams));
        //添加自定义的静态函数，制作扰乱的工具函数：
        List<IMethod.StaticMethod> newStaticMethods = randomStaticMethods(argumentParams.getBasePackage());
        //修改两种共有的函数
        Collection<String> methodSignatures = methodDeclarationMap.keySet();
        if (argumentParams.isReferenceMethodOnly()) {
            methodSignatures = CollectionUtil.intersection(methodDeclarationMap.keySet(), methodCallExprMap.keySet());
        }
        for (String methodSignature : methodSignatures) {
            MethodDeclaration methodDeclaration = methodDeclarationMap.get(methodSignature);
            //1.给当前的类添加一个loggerPrint函数
            methodDeclaration.findAncestor(ClassOrInterfaceDeclaration.class).ifPresent(p -> {
                List<MethodDeclaration> methods = p.getMethodsByName(RandomExpression.getLoggerPrintName());
                if (methods.isEmpty()) {
                    MethodDeclaration m = p.addMethod(RandomExpression.getLoggerPrintName(), Modifier.Keyword.PUBLIC, Modifier.Keyword.STATIC);
                    m.addParameter(new Parameter(StaticJavaParser.parseType("Object"), "values").setVarArgs(true));
//                    String body = "{try {logger.finest(\"" + RandomExpression.getName() + "\"+Arrays.toString(values));}catch (Exception ignored){}}";
                    String body = IMethod.getLoggerPrintBody(config);
                    m.setBody(StaticJavaParser.parseBlock(body));
                }
            });
            //2.随机获取一个函数, 对函数添加参数
            List<IMethod.AddTypeInfo> params = new ArrayList<>();
            IMethod.StaticMethod staticMethod = newStaticMethods.get(random.nextInt(newStaticMethods.size()));
            NodeList<Parameter> staticParams = staticMethod.md.getParameters();
            NodeList<Parameter> parameters = methodDeclaration.getParameters();
            boolean lastParamIsVarArg = !parameters.isEmpty() && parameters.get(parameters.size() - 1).isVarArgs();
            int paramNum = Math.max(random.nextInt(argumentParams.getMaxAddParamNum()), argumentParams.getMinAddParamNum());
            for (int i = 0; i < Math.max(paramNum, staticParams.size()); i++) {
                Type paramsType = i < staticParams.size() ? staticParams.get(i).getType() : RandomExpression.randomType();
                String paramName = RandomExpression.getName() + "_" + i;
                List<String> curParams = parameters.stream().map(Parameter::getNameAsString).collect(Collectors.toList());
                while (curParams.contains(paramName)) {
                    paramName = RandomExpression.getName() + i;
                }
                int insertIndex = random.nextInt(parameters.size() + (lastParamIsVarArg ? 0 : 1));
                params.add(new IMethod.AddTypeInfo(insertIndex, paramsType, paramName));
                parameters.add(insertIndex, new Parameter(paramsType, paramName));
            }
            methodDeclaration.setParameters(parameters);
            //3.添加上面随机获取函数导入类
            methodDeclaration.findCompilationUnit().ifPresent(tcu -> {
//                tcu.addImport("java.util.logging.Logger");
                tcu.addImport(IMethod.getLoggerImport(config));
                tcu.addImport("java.util.Arrays");
                tcu.addImport(staticMethod.packageName + "." + staticMethod.className);

                if (config.getRunningStatus() == RunningEnum.JAVA_PROJECT) {
                    // 获取类定义
                    tcu.getPrimaryType()
                            .flatMap(pt -> pt.toClassOrInterfaceDeclaration())
                            .ifPresent(clazz -> {
                                // 检查是否已有 logger 字段
                                boolean hasLogger = clazz.getFields().stream()
                                        .flatMap(fd -> fd.getVariables().stream())
                                        .anyMatch(var -> var.getNameAsString().equals("logger"));
                                if (!hasLogger) {
                                    // 构造 static final Logger 字段
                                    String className = clazz.getNameAsString();
                                    FieldDeclaration loggerField = StaticJavaParser.parseBodyDeclaration(
                                            "private static final Logger logger = Logger.getLogger(" + className + ".class.getName());"
                                    ).asFieldDeclaration();
                                    clazz.addMember(loggerField);
                                }
                            });
                }
            });
            //4.对函数添加与主逻辑无关的代码片段，用于扰乱代码逻辑
            if (!params.isEmpty()) {
                //添加无效的扰乱代码
                methodDeclaration.getBody().ifPresent(blockStmt -> {
                    Map<String, String> paramInfos = params.stream().collect(Collectors.toMap(key -> key.paramName, value -> value.paramsType.asString()));
                    List<Statement> addStatements = RandomExpression.generateStatement(paramInfos);
                    NodeList<Statement> thisStatements = blockStmt.getStatements();
                    int addIndex = -1;
                    for (Statement statement : addStatements) {
                        int insertIndex = thisStatements.isEmpty() ? 0 : random.nextInt(thisStatements.size());
                        addIndex = Math.max(insertIndex, addIndex + 1);
                        thisStatements.add(addIndex, statement);
                    }
                    //调用工具类中的无效函数
                    int insertIndex = thisStatements.isEmpty() ? 0 : random.nextInt(thisStatements.size());
                    String returnName = RandomExpression.getName() + RandomExpression.getName();
                    String staticMethodCode = staticMethod.className + "." + staticMethod.md.getName() + "(";
                    List<String> thisParamNames = new ArrayList<>();
                    for (int num = 0; num < staticParams.size(); num++) {
                        thisParamNames.add(params.get(num).paramName);
                    }
                    staticMethodCode = staticMethodCode + String.join(", ", thisParamNames) + ");";
                    staticMethodCode = "Object " + returnName + " = " + staticMethodCode;
                    thisStatements.add(insertIndex, StaticJavaParser.parseStatement(staticMethodCode));
                    String returnPrint = RandomExpression.println + "(\"fun" + returnName + "\", " + returnName + ");";
                    thisStatements.add(insertIndex + 1, StaticJavaParser.parseStatement(returnPrint));
                    blockStmt.setStatements(thisStatements);
                });
            }
            //5.对引用点添加参数信息
            List<MethodCallExpr> methodCallExprList = methodCallExprMap.get(methodSignature);
            if (null != methodCallExprList) {
                for (MethodCallExpr mc : methodCallExprList) {
                    NodeList<Expression> arguments = mc.getArguments();
                    for (IMethod.AddTypeInfo aty : params) {
                        arguments.add(aty.insertIndex, RandomExpression.generateValue(aty.paramsType.asString()));
                    }
                    mc.setArguments(arguments);
                }
            }
        }
        //6.保存修改的文件
        ObfuscatorUtils.saveObfuscatedFiles(fileMap);
    }

    //是否启动这个功能
    @Override
    public boolean supports(ObfuscatorConfig config) {
        return ObjectUtil.isNotEmpty(config.getMethodArgumentParams());
    }


    /**
     * 遍历出哪些函数可以进行参数修改
     **/
    @Override
    public void visit(ClassOrInterfaceDeclaration cid, MethodArgumentParams config) {
        super.visit(cid, config);
        //混淆的函数不是接口与内部类
        if (!cid.isInnerClass() && !cid.isInterface()) {
            cid.getParentNode().ifPresent(parentNode -> {
                if (parentNode instanceof CompilationUnit) {
                    CompilationUnit cu = (CompilationUnit) parentNode;
                    cu.getPackageDeclaration().ifPresent(packageDeclaration -> {
                        String packageName = packageDeclaration.getNameAsString();
                        if (config.getPackageList().isEmpty() || config.getPackageList().stream().anyMatch(packageName::startsWith)) {
                            for (MethodDeclaration md : cid.getMethods()) {
                                //==============================⬇️================================//
                                //新加方法:防止抽象类重写父类接口方法没加override注解仍混淆的bug
                                //解析方法
                                ResolvedMethodDeclaration rmd = md.resolve();
                                //判断是不是接口或父类里已有的方法
                                boolean overrides = false;
                                try {
                                    overrides = rmd.declaringType()
                                            .getAllAncestors().stream()
                                            .anyMatch(ancestor -> {
                                                try {
                                                    return ancestor.getTypeDeclaration().isPresent() &&
                                                            ancestor.getTypeDeclaration().get()
                                                                    .getDeclaredMethods().stream()
                                                                    .anyMatch(am -> am.getSignature().equals(rmd.getSignature()));
                                                } catch (Exception e) {
                                                    return false;//出现异常就不算重写方法
                                                }
                                            });
                                } catch (Exception e) {
                                    overrides = false; //出现异常就不算重写方法
                                }
                                if (!overrides &&
                                        //==============================⬆️================================//
                                        md.getAnnotationByName("Override").isEmpty() &&
                                        md.getParameters().size() < 10 &&
                                        !md.isNative() && !md.isAbstract() &&
                                        !RandomExpression.println.equalsIgnoreCase(md.getNameAsString())
                                ) {
                                    try {
                                        //第一版判断
                                        String methodSignature = md.resolve().getQualifiedSignature();
                                        methodDeclarationMap.put(methodSignature, md);
                                        //==============================⬇️================================//
                                        //新版判断:只将属于指定包下的普通方法放入集合中
//                                        ResolvedMethodDeclaration rmd = md.resolve();
//                                        String methodPkg = rmd.declaringType().getPackageName();
//                                        if (config.getPackageList().stream().anyMatch(methodPkg::startsWith)) {
//                                            methodDeclarationMap.put(rmd.getQualifiedSignature(), md);
//                                        }
                                        //==============================⬆️================================//
                                    } catch (Exception ignore) {
                                    }
                                }
                            }
                        }
                    });
                }
            });
        }
    }

    /**
     * 这里处理函数调用时，对参数的修改
     **/
    @Override
    public void visit(MethodCallExpr mc, MethodArgumentParams config) {
        super.visit(mc, config);
        mc.toMethodCallExpr().ifPresent(resolvedMethod -> {
            try {
                //第一版判断
//                String methodSignature = resolvedMethod.resolve().getQualifiedSignature();
//                methodCallExprMap.computeIfAbsent(methodSignature, k -> new ArrayList<>()).add(mc);
                //==============================⬇️================================//
                //第二版判断:限定调用方只统计属于指定包下的方法
                ResolvedMethodDeclaration rmd = resolvedMethod.resolve();
                if (rmd.getPackageName() != null &&
                        config.getPackageList().stream().anyMatch(rmd.getPackageName()::startsWith)) {
                    methodCallExprMap.computeIfAbsent(rmd.getQualifiedSignature(), k -> new ArrayList<>()).add(mc);
                }
                //==============================⬆️================================//
            } catch (Exception ignored) {
                if (print) {
                    System.err.println(ignored.getMessage());
                }
            }
        });
    }

    /**
     * 随机生成工具类
     **/
    public List<IMethod.StaticMethod> randomStaticMethods(String basePackage) {
        //当前新增的包
        List<IMethod.StaticMethod> newStaticMethods = new ArrayList<>();
        List<String> packages = new ArrayList<>();
        //原本是5
        for (int i = 0; i < IMethod.INVALID_PACKAGE_NUM; i++) {
            packages.add("z" + RandomUtil.randomString(RandomUtil.BASE_CHAR, 5));
        }
        //遍历静态方法，并保存方法到不同的文件中
        List<Integer> useList = new ArrayList<>();
        while (useList.size() != tempStaticMethods.size()) {
            //对方法进行随机分组
            List<MethodDeclaration> part = new ArrayList<>();
            for (int j = 0; j < random.nextInt(3) + 3; j++) {
                if (useList.size() == tempStaticMethods.size()) {
                    continue;
                }
                int index = random.nextInt(tempStaticMethods.size());
                while (useList.contains(index)) {
                    index = random.nextInt(tempStaticMethods.size());
                }
                MethodDeclaration thisMd = tempStaticMethods.get(index);
                if (random.nextInt(10) > 5) {
                    thisMd.getName().setIdentifier(thisMd.getName().asString() + RandomExpression.getName().toLowerCase());
                } else if (random.nextInt(10) > 5) {
                    thisMd.getName().setIdentifier(RandomExpression.getName().toLowerCase() + thisMd.getName().asString());
                }
                part.add(thisMd);
                useList.add(index);
            }
            //新增文件并添加函数到脚本中
            if (!part.isEmpty()) {
                String thisPackage = basePackage + "." + packages.get(random.nextInt(packages.size()));
                String className = "UT" + RandomExpression.getName().toLowerCase();
                CompilationUnit cu = new CompilationUnit();
                cu.setPackageDeclaration(thisPackage);
                cu.addClass(className).setPublic(true).setFinal(true);
                cu.getClassByName(className).ifPresent(cid -> {
                    for (MethodDeclaration item : part) {
                        cid.addMember(item);
                    }
                });
                try {
                    Path savePath = Paths.get(this.sourceRoot, thisPackage.replace(".", "/"));
                    if (!savePath.toFile().exists()) {
                        if (!savePath.toFile().mkdirs()) {
                            throw new RuntimeException("mkdirs error");
                        }
                    }
                    Files.writeString(Paths.get(savePath.toString(), className + ".java"), cu.toString());
                    for (MethodDeclaration item : part) {
                        newStaticMethods.add(new IMethod.StaticMethod(thisPackage, className, item));
                    }
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        }
        return newStaticMethods;
    }
}
